import collections
import datetime
import functools
import logging
import os
import traceback

import pandas as pd
from absl import app, flags
from tornado.ioloop import IOLoop, PeriodicCallback

from coin.base.timestamp import get_timestamp
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.support.proto_log.logic.util import StratInfo, run_from_strat_log_archive
from coin.tool.strat_monitor.util.monitor_util import (
    get_strat_from_json_config,
    get_topic_from_strat)
from coin.tool.strat_monitor.handlers.strategy_alarm_handler import (
    RateLimitAlarm, StrategyAlarmHandler, TopicNotFoundAlarm)
from coin.util.queue.config import KafkaConfig
from coin.util.queue.tools.kafka_archive import run_from_kafka_archive
from coin.util.queue.tools.kafka_topic import parse_kafka_topic

from coin.base.param_util import to_list
from xunkemgmt_client.client.util.query_util import (
    query_accounts,
    query_strategies,
    query_applications)
from xunkemgmt_client.tool.slack_noti import send_to_slack


FLAGS = flags.FLAGS


IGNORE_ALARMS = [RateLimitAlarm]
NOT_MENTION_ALARMS = [RateLimitAlarm]


def _force_quit(func):
  def wrapper(*args, **kwargs):
    try:
      func(*args, **kwargs)
    except Exception:
      logging.info(traceback.format_exc())
      os._exit(1)
  return wrapper


def report_slack(strategy_meta, alarm_info, channel, timestamp, mention=True):
  """
  alarm_info = {
    'lm_agg2_binance_spot_btcquote_1': [
      "Rate limit hit: ('Spot', 'Binance', 'xunke49'): ['CANCEL', 'PLACE']"
      "High order error ratio in last 5 min: {'ETC-USDT.PERPETUAL': {'': 'Margin is insufficient.'}}"
    ],
    ...
  }
  """
  strategy_info = []
  for strat_name, alarm_msgs in alarm_info.items():
    try:
      trader = strategy_meta.set_index('strategy_name').loc[strat_name, 'trader']
      strategy_group = \
          strategy_meta.set_index('strategy_name').loc[strat_name, 'strategy_group']
      account_info_list = strategy_meta.set_index('strategy_name').loc[strat_name, 'account_info']
      mea_list = [f"{account_info['market_type']}.{account_info['exchange']}.{account_info['api_version']}"
        for account_info in account_info_list]
      owner_list = [account_info['owner'] for account_info in account_info_list]
      mea = ','.join(mea_list)
      owner = ','.join(owner_list)
    except Exception:
      trader = strategy_group = None
      logging.info('strategy name not found %s' % strat_name)
    strategy_info.append((strat_name, strategy_group, trader, mea, owner))
  strategy_df = pd.DataFrame(strategy_info, columns=['strategy_name', 'strategy_group', 'trader', 'mea', 'owner'])

  traders = strategy_df['trader'].drop_duplicates().tolist()
  mention_target = [trader if trader is not None else 'chensili' for trader in traders]
  strategy_df.set_index(['trader', 'strategy_group'], inplace=True)
  msg = "```"
  utc_dt = datetime.datetime.utcfromtimestamp(timestamp / 1e+9).strftime('%Y-%m-%d %H:%M:%S')
  for index in set(strategy_df.index.tolist()):
    result = strategy_df.loc[[index]]
    trader, group = index
    msg += "strategy_group: %s, trader: %s, #abnormal strats: %s (%s)" % (group, trader,
                                                                          len(result), utc_dt)
    for _, strat in result.iterrows():
      alarm_msgs = alarm_info.get(strat['strategy_name'], [])
      if len(alarm_msgs) == 0:
        msg += "\n    %s" % strat['strategy_name']
      else:
        msg += "\n    %s, %s, %s: %s" % (strat['strategy_name'], strat['mea'], strat['owner'], ',  '.join(alarm_msgs))
    msg += "\n \n"
  msg += "```"
  mention_list = mention_target if mention else None
  send_to_slack(msg, channel, 'msg', mention_list=mention_list)


class StrategyLogAlarm(object):
  def __init__(self,
               *,
               strat_infos,
               current_ts=None,
               business_unit=None,
               ignored_strategy_filename=None,
               kafka_config_filename=None,
               strat_log_root_dir=None,
               logger=None):
    self._strat_infos = strat_infos
    self._strat_list = list(set(
        [strat_info.strategy_name for strat_info in strat_infos]))
    self._business_unit = business_unit
    self._ignored_strategy_filename = ignored_strategy_filename
    self._kafka_config_filename = kafka_config_filename
    self._kafka_config = None
    self._strat_log_root_dir = strat_log_root_dir
    if kafka_config_filename is not None:
      self._kafka_config = KafkaConfig.from_cmd_config(kafka_config_filename)
    self._history_all_alarm = collections.defaultdict(dict)
    self._end_ts = current_ts
    self._warm_up = False
    self._handlers = {}
    self._strategy_universe = None
    self._logger = logger or logging.getLogger(__name__)
    self.update_handler_info()

  def register_strategy_handler(self):
    self._logger.info("Register handler")
    if self._strategy_universe is None:
      self._logger.info("Waiting for query strategy universe")
      return
    for strategy_name in self._strat_list:
      if self._handlers.get(strategy_name) is None:
        strategy_info = self._strategy_universe.loc[strategy_name].dropna().to_dict()
        if strategy_info.get('strategy_name') is None or strategy_info.get('active') is None:
          self._logger.info("Failed to generate handler for %s" % strategy_name)
          continue
        self._handlers[strategy_name] = StrategyAlarmHandler(strategy_info=strategy_info, logger=self._logger)

  def on_log(self, strategy_name, machine, ts, log):
    pb = StrategyLog()
    pb.ParseFromString(log)
    self._handlers[strategy_name].handle(ts, machine, pb)
    return

  def on_kafka_log(self, partition, record):
    _, strat_info = parse_kafka_topic(record.topic)
    self.on_log(strat_info.strategy_name, strat_info.machine,
                record.timestamp * 1e6, record.value)

  def query_meta_info(self):
    account_df = query_accounts()
    application_df = query_applications()
    strategy_df = query_strategies()
    if self._business_unit is not None:
      business_unit = to_list(self._business_unit)
      strategy_df = strategy_df[strategy_df['business_unit'].isin(business_unit)]

    def _agg_account_info(strategy_id):
      account_ids = application_df['account_id'][application_df['strategy_id'] == strategy_id].tolist()
      accunt_info = account_df[account_df['account_id'].isin(account_ids)].to_dict('record')
      return accunt_info
    strategy_df['account_info'] = strategy_df['strategy_id'].apply(_agg_account_info)
    strategy_df.index = strategy_df['strategy_name']
    self._strategy_universe = strategy_df

  def update_strategy_info(self):
    for strategy_name, handler in self._handlers.items():
      if strategy_name not in self._strategy_universe.index:
        # In spite of low chance, it is still possible.
        # Eg. account manager registers a new strategy in meta db and removes some time
        # later realizing that strategy info is wrong.
        handler.update_strategy_info({'active': False})
      else:
        strategy_info = self._strategy_universe.loc[strategy_name].dropna().to_dict()
        handler.update_strategy_info(strategy_info)

  @_force_quit
  def update_handler_info(self):
    self.query_meta_info()
    self.register_strategy_handler()
    self.update_strategy_info()

  def update_all_timestamp(self, ts):
    for strategy, handler in self._handlers.items():
      handler.update_current_ts(ts)

  def check_no_topic_strategy(self):
    topic_not_found_map = {}
    if self._kafka_config_filename is None:
      return {}
    topics = get_topic_from_strat(None, self._kafka_config, active=True)
    strats_with_topic = [topic.split('_', 2)[-1] for topic in topics]
    active_strats = self._strategy_universe[self._strategy_universe['active'] == True]
    strats_without_topic = [strat for strat in active_strats['strategy_name'].drop_duplicates() if strat not in strats_with_topic]
    if strats_without_topic:
      topic_not_found_map.update({strat: TopicNotFoundAlarm.alarm_msg() for strat in strats_without_topic})
      return {TopicNotFoundAlarm: topic_not_found_map}
    else:
      return {}

  @_force_quit
  def _reset(self):
    self._logger.debug("Reset")
    self._history_all_alarm.clear()

  def agg_all_alarms(self):
    all_alarm_total = {}
    for strategy, handler in self._handlers.items():
      try:
        alarm_msgs = handler.get_all_alarms()
      except Exception as e:
        self._logger.info("%s alarm collect failed." % strategy)
        logging.info(traceback.format_exc())
        alarm_msgs = {}
      for alarm_type, alarm_msg in alarm_msgs.items():
        alarm_summary = all_alarm_total.get(alarm_type, {})
        alarm_summary.update(alarm_msg)
        all_alarm_total.update({alarm_type: alarm_summary})
    no_topic_alarm = self.check_no_topic_strategy()
    if no_topic_alarm:
      all_alarm_total.update(no_topic_alarm)
    return all_alarm_total

  def check_alarms(self):
    self._all_alarm_info = self.agg_all_alarms()
    ignored_strats = []
    try:
      if self._ignored_strategy_filename is not None:
        ignored_strats = set(get_strat_from_json_config(self._ignored_strategy_filename))
    except:
      logging.error('Fail to get ignored strategies.')
      logging.error(traceback.format_exc())
    all_alarm_infos = collections.defaultdict(dict)
    ignored_alarm_infos = collections.defaultdict(dict)
    new_alarm_infos = collections.defaultdict(dict)
    for alarm_type, alarm_info in self._all_alarm_info.items():
      history_alarm_info = self._history_all_alarm[alarm_type]
      for strat_name, alarm_msg in alarm_info.items():
        if alarm_type in IGNORE_ALARMS or strat_name in ignored_strats:
          ignored_alarm_infos[alarm_type][strat_name] = alarm_msg
          continue
        all_alarm_infos[alarm_type][strat_name] = alarm_msg
        if strat_name not in history_alarm_info:
          new_alarm_infos[alarm_type][strat_name] = alarm_msg
      self._history_all_alarm[alarm_type] = alarm_info
    if len(all_alarm_infos) > 0:
      self._logger.info("All alarms msgs:")
      self._logger.info(all_alarm_infos)
    if len(ignored_alarm_infos) > 0:
      self._logger.info("Ignored alarms msgs:")
      self._logger.info(ignored_alarm_infos)
    if len(new_alarm_infos) > 0:
      self._logger.info("New alarms msgs:")
      self._logger.info(new_alarm_infos)
    return new_alarm_infos

  def _convert_alarm_info_per_trader(self, alarm_infos):
    converted_alarm_infos = {}
    strategy_df = self._strategy_universe
    for alarm_type, alarm_info in alarm_infos.items():
      for strat, msg in alarm_info.items():
        try:
          trader = strategy_df.set_index('strategy_name').loc[strat, 'trader']
        except Exception:
          trader = None
        alarm_info_per_trader = converted_alarm_infos.get(trader, collections.defaultdict(list))
        alarm_info_per_trader[strat].append(msg)
        converted_alarm_infos.update({trader: alarm_info_per_trader})
    return converted_alarm_infos

  @_force_quit
  def report_abnormals(self, channel=None):
    self._logger.info("collect alarm msgs")
    new_alarm_info = self.check_alarms()
    channel = channel or FLAGS.slack_receiver
    if len(new_alarm_info) == 0:
      return
    if channel is None:
      self._logger.info("No channel to send")
      return
    alarm_infos = {
        alarm_type: alarm_info for alarm_type, alarm_info in new_alarm_info.items()
        if alarm_type not in NOT_MENTION_ALARMS
    }
    alarm_infos_not_mention = {
        alarm_type: alarm_info for alarm_type, alarm_info in new_alarm_info.items()
        if alarm_type in NOT_MENTION_ALARMS
    }
    alarm_infos_per_trader = self._convert_alarm_info_per_trader(alarm_infos)
    alarm_infos_not_mention_per_trader = self._convert_alarm_info_per_trader(
        alarm_infos_not_mention)

    for trader, alarm_info in alarm_infos_per_trader.items():
      report_slack(self._strategy_universe, alarm_info, channel, self._end_ts)
    for trader, alarm_info in alarm_infos_not_mention_per_trader.items():
      report_slack(self._strategy_universe, alarm_info, channel, self._end_ts, mention=False)

  def warm_up(self):
    self._warm_up = True
    end_ts = self._end_ts - int(datetime.timedelta(minutes=5).total_seconds() * 1e+9)
    start_ts = end_ts - int(datetime.timedelta(minutes=60).total_seconds() * 1e+9)
    self.read_log(start_ts, end_ts)
    self._warm_up = False
    self._logger.info("Warm up end")

  @_force_quit
  def read_log(self, start_ts=None, end_ts=None):
    start_ts = start_ts or self._end_ts
    end_ts = end_ts or get_timestamp()
    self._logger.debug("Read log from %s to %s" % (datetime.datetime.utcfromtimestamp(
        start_ts / 1e+9), datetime.datetime.utcfromtimestamp(end_ts / 1e+9)))
    if self._kafka_config is not None:
      topics = get_topic_from_strat(self._strat_list, self._kafka_config)
      run_from_kafka_archive(topics=topics,
                             kafka_config=self._kafka_config,
                             timestamp_from=int(start_ts),
                             timestamp_to=int(end_ts),
                             callback=self.on_kafka_log)
    else:
      assert len(self._strat_infos) == 1
      machine = self._strat_infos[0].machine
      strategy_name = self._strat_infos[0].strategy_name
      run_from_strat_log_archive(
          on_log_callback=functools.partial(self.on_log, strategy_name, machine),
          start_time=datetime.datetime.utcfromtimestamp(start_ts/10**9),
          end_time=datetime.datetime.utcfromtimestamp(end_ts/10**9),
          root_dir=self._strat_log_root_dir,
          machine=machine,
          strategy_name=strategy_name)
    self._logger.debug("Read log end")
    self._end_ts = end_ts
    self.update_all_timestamp(end_ts)

  def start(self, alarm_frequency):
    self.warm_up()
    PeriodicCallback(self.update_handler_info, 300 * 1000).start()
    PeriodicCallback(self.read_log, alarm_frequency / 2 * 1000).start()
    PeriodicCallback(self.report_abnormals, alarm_frequency * 1000).start()
    PeriodicCallback(self._reset, 1800 * 1000).start()

def main(_):
  ioloop = IOLoop.current()
  # TODO(chensili): remove these. only query in StrategyLogAlarm. no need here
  strategy_universe = query_strategies(as_proto=True)
  filtered = strategy_universe
  business_unit = FLAGS.business_unit
  if FLAGS.business_unit is not None:
    business_unit = [elem.strip() for elem in FLAGS.business_unit.split(',')]
    filtered = [elem for elem in filtered if elem.business_unit in business_unit]
  strat_list = [elem.strategy_name for elem in filtered]
  strat_infos = [StratInfo(strategy_name=strat, machine=None, trading_date=None) for
                 strat in strat_list]
  alarm = StrategyLogAlarm(strat_infos=strat_infos,
                           ignored_strategy_filename=FLAGS.ignored_strategy_filename,
                           current_ts=get_timestamp(),
                           business_unit=business_unit,
                           kafka_config_filename=FLAGS.kafka_config_filename)
  alarm.start(alarm_frequency=FLAGS.alarm_frequency)
  if FLAGS.exit_after_min:
    ioloop.add_timeout(datetime.timedelta(minutes=FLAGS.exit_after_min), ioloop.stop)

  try:
    ioloop.start()
  except (KeyboardInterrupt):
    logging.info('KeyboardInterrupt')
    return


if __name__ == '__main__':
  logging.basicConfig(level='INFO', format='%(levelname)8s %(asctime)s %(message)s')
  flags.DEFINE_string('ignored_strategy_filename',
                      '../../coin_deploy/support_monitor/config/service_mode_strategy_list.json',
                      'strategy to be ignored')
  flags.DEFINE_string('kafka_config_filename',
                      '../../coin_deploy/support_monitor/config/kafka_aws_config.json',
                      'kafka config')
  flags.DEFINE_string('business_unit', None, 'filter by business_unit')
  flags.DEFINE_string('slack_receiver', None, 'channel to receive alarm msg')
  flags.DEFINE_integer('alarm_frequency', 120, 'alarm frequency')
  flags.DEFINE_integer('exit_after_min', 240, 'exit after minute')
  app.run(main)
